﻿using FlyIL;
using FlyIL.Advance;
using FlyIL.Emu;
using FlyIL.Lifter;
using FlyIL.Misc;
using System;
using System.Collections.Generic;
using System.Linq;
using static FlyIL.AssemblerRegisters;
using static FlyIL.OperandMaker;

namespace FuckSp;

public class SpEntry
{


    /********************************************************
     * 
     * 虚拟代码（或者说  next ）地址 由 esi 指定 
     * eax 用来 计算 下一处 地址 
     * 
     * sp 不存在 vminit的 概念 (不会平衡栈)， 因此 做为一个整段汇编 ， 进行化简即可 ， hook 同理 
     * ? sp的花指令 有否 跨 handler 机制？
     * - 将sp的花指令拼接并化简 
     *      - esi中取常量进行直接替换
     *      - 对 vJmp 需要进行额外收集 
     * - 遇到 vcall ， 添加一个新的 VmEntry , 并收集 
     * - 
     **************************************/



    readonly Detector Detector;


    readonly List<Operand> remove1 = new() { _regs_[R1D], _regs_[R6D], _regs_[R4D], _regs_[CF] };
    readonly List<Operand> remove2 = new() { _regs_[R1D], _regs_[R6D], _regs_[R4D], _regs_[CF] ,
        _regs_[R3D], _regs_[R7D], _regs_[R4D], _regs_[EFlag] };



    public SpEntry(Detector detector)
    {
        Detector = detector;


    }



    /// <summary>
    /// 先将 handler 化简, 再将它复制到 主context 
    /// </summary>
    /// <param name="address"></param>
    /// <param name="SpContext"></param>
    /// <param name="handler"></param>
    /// <param name="_stop"></param>
    /// <param name="removes"></param>
    /// <param name="_start"></param>
    void CollectHandler(ulong address, Context SpContext, Context handler, Inst _stop, List<Operand> removes, Inst _start = null)
    {
        uint subip = 0;
        handler.SimplifyBetween(handler.Entry, _stop);
        /// 移除对 eax ，esi 的计算代码
        handler.RemoveDeadBatch(_stop, removes);

        for (var node = handler.Entry; node != null && node != _stop; node = handler.GetNextIfExist(node))
        {
            if (node.Opt == Opt.nop)
                continue;
            if (node.IsJmp)
                continue;

            SpContext.Make(address, subip).CopyFrom(node);
            subip += 2;

            if (subip == 0)
                SpContext.Last.Note = $"@ {handler.Entry.Vip.ToString("X")}";

        }


    }



    /// <summary>
    /// 在运行中 ，修正 即时数 以及 取址偏移
    /// </summary>
    void FixIndexOrImm(Cpu cpu, Context handler, Inst node)
    {

        if (node.HavePrefix)
            return;

        if (node.OpLast.IsMatch(_any_[Any + R1D]))
        {
            var op = node.OpLast;
            op.FixIndex((uint)cpu.Context.RAX);
            node.OpLast = op;

            var eax = _regs_[R1D];
            handler.RemoveDeadCode(node, ref eax);
        }

        if (node.O1.IsMatch(_any_[Any + R1D]))
        {
            var op = node.O1;
            op.FixIndex((uint)cpu.Context.RAX);
            node.O1 = op;

            var eax = _regs_[R1D];
            handler.RemoveDeadCode(node, ref eax);
        }

        /// 从 [esi+???]  或者从 esi 赋值后的 寄存器  取常量  
        if (node.OpLast.IsPurePtr(RegCode.r6d, null, null) || (node.OpLast.IsPtr && cpu.AreRegEqual(node.OpLast.Base, RegCode.r6d)))
        {
            uint offset = node.OpLast.u32;
            int byte_size = node.OpLast.ByteSize;

            ulong ul = (ulong)(byte_size == 4 ? cpu.Reader.Read32(cpu.Context.RSI + offset)
                : byte_size == 2 ? cpu.Reader.Read16(cpu.Context.RSI + offset)
                : cpu.Reader.Read8(cpu.Context.RSI + offset));

            node.OpLast = Operand.Make(ul, byte_size);
        }
    }


    ulong FindReadNext(Context handler)
    {

        //if (handler.EntryPoint == 0x0425BDAE)
        //    ;

        for (var node = handler.Last; node != null; node = handler.GetPreviousIfExist(node))
            /// 从 [esi+???]  或者从 esi 转移后的 寄存器  取常量  
            if (node.OpLast == _dword_[R6D])
            {
                return node.Vip;
            }

        throw new NotImplementedException();
    }


    void _init_checkVsp()
    {

        if (Detector._checkVsp != null)
            return;


        var con = Detector.Vcontext;

        Detector._checkVsp = new(0x2000);
        Tag b = new("b", 0);

        con.Label(Detector._checkVsp).mov(R3D, _dword_[Esp + 4]);
        con.Make().mov(R1D, _dword_[Esp + 8]);
        con.Make().mov(R4D, R7D);
        con.Make().add(R7D, 0x14c);
        con.Make().sub(R3D, R4D);

        con.Make().cmp(R3D, R1D);

        con.Make().jge(b);

        con.Make().mov(_dword_[R7D + 0x88], R1D);
        con.Make().mov(R4D, Esp);
        con.Make().mov(_dword_[R7D + 0x8c], R4D);
        con.Make().mov(R4D, R7D);
        con.Make().add(R4D, 0x14c);

        con.Make().sub(R4D, _dword_[R7D + 0x8c]);
        con.Make().mov(_dword_[R7D + 0x90], R4D);
        con.Make().mov(R3D, _dword_[R7D + 0x88]);
        con.Make().sub(Esp, R3D);

        con.Make().mov(R3D, _dword_[R7D + 0x90]);
        //con.Make().mov(R6D, _dword_[R7D + 0x8c]);
        con.Make().mov(R4D, R7D);
        con.Make().mov(R7D, Esp);

        con.Make()._rep().mov(_byte_[R7D], _byte_[R6D]);
        con.Make()._if(ExprRegister.DF == 0).add(R7D, 1);
        con.Make()._if(ExprRegister.DF == 0).add(R6D, 1);
        con.Make()._if(ExprRegister.DF == 1).sub(R7D, 1);
        con.Make()._end()._if(ExprRegister.DF == 1).sub(R6D, 1);

        con.Make().sub(R4D, R1D);
        con.Make().mov(R7D, R4D);

        con.Label(b).ret(0x8);

    }


    Context prepareHandler(bool is64, ulong addr, IReader reader)
    {
        /// todo: 首个片段 使用 checkRet , 会造成栈错位 ，禁用后无此问题
        var handler = new Context(is64, addr);

        /// 屏蔽这些 计算 
        handler.OptimizerOption.IgnoreOpts.Add(Opt.adc);
        var lifter = new Lifters.Intel(is64, reader);
        var config = new Lifters.IntelOption { CheckRet = false, CallAsjmp = true, TrustJcc = true };

        lifter.Collect(handler, config, addr);

        ulong read_next = FindReadNext(handler);
        handler.Marks["next"] = read_next;

        var read_node = handler[read_next];
        /// 如果是 vjcc 需要把 未跳转的代码 也收集过来
        if (Detector.CheckHandle(handler, out Inst jno) >= HandleType.jcc)
        {
            lifter.CollectNext(handler, config, jno.Vip);
        }
        else
        {
            /// 没有条件跳转 ， 移除 cmp test 语句
            handler.NopFlow(handler.Entry, read_node, Opt.cmp, Opt.test);

        }

        /**************************
         * 在收集后 即对 handler 化简 ，即可验证化简的正确性 ，又可提升效率 
         * ***********************/
        //if (handler.EntryPoint == 0x00000000042e5606)
        //    ;

        handler.SimplifyBetween(handler.Entry, read_node);

        /// 移除这些操作数 不会影响 cpu 的执行
        handler.RemoveDeadBatch(handler.GetPreviousIfExist(read_node), new() { _regs_[CF], _regs_[R1D] });
        handler.ClearNop();


        Detector.handlers[addr] = handler.Pack();

        return handler;
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="is64"></param>
    /// <param name="reader"></param>
    /// <param name="report"></param>
    /// <param name="hAddr"> handler 地址 </param>
    /// <param name="vip">入口的esi值，即虚拟地址</param>
    public void Collect(bool is64, IReader reader, IReport report, ulong hAddr, ulong vip = 0)
    {

        // 收集到总的上下文中  
        var con = Detector.Vcontext;

        Cpu cpu = new(is64, reader);
        cpu.SetReg64(RegCode.r6, vip);

        ulong handle_addr = hAddr;
        ulong cip = 0;

        bool removeMore = false;

        for (int i = 0; ; handle_addr = cpu.NextIp, i++)
        {
            cip = cpu.GetReg64(RegCode.r6);


            //if (cip == 0x042E69C6)
            //    ;

            Context handler;
            char firstFound = '+';
            string text = null;

            if (!Detector.handlers.ContainsKey(handle_addr))
            {
                handler = prepareHandler(is64, handle_addr, reader);
            }
            else
            {
                firstFound = ' ';
                handler = Context.UnPack(Detector.handlers[handle_addr]);
            }


            var h_type = Detector.CheckHandle(handler, out Inst jnode);
            if (h_type == HandleType.checkVsp)
            {
                ;
                handler.Dump();
                _init_checkVsp();
                cpu.Run(handler, handler.Entry, null);
                con.Make(cip).call(Detector._checkVsp);

                text = string.Format("第{0:d4}段, {1} VIP= {2:X8} , HANDLE= {3:X8} , {4}", i, firstFound, cip, handle_addr, h_type.ToString());
                Console.WriteLine(text);

                continue;


            }
            else if (h_type == HandleType.jcc)
            {

                handler.Dump();
                /// 使用递归  继续 收集 
                Collect(is64, reader, report, 111);
                return;
            }
            else if (h_type == HandleType.jmp)
            {
                ;
                cpu.Run(handler, handler.Entry, null);
                con.Make(cip).jmp(cpu.GetReg64(RegCode.r6));

                text = string.Format("第{0:d4}段, {1} VIP= {2:X8} , HANDLE= {3:X8} , {4}", i, firstFound, cip, handle_addr, h_type.ToString());
                Console.WriteLine(text);

                continue;
            }

            Inst _read_next = handler[(ulong)handler.Marks["next"]];
            cpu.Run(handler, handler.Entry, null, FixIndexOrImm);

            /// 如果是 vinit ,那么修正 cip 地址为 下一个  cip -4  
            if (cip == 0)
                cip = cpu.GetReg64(RegCode.r6) - 4;



            // 当前这一串代码的 vip 
            text = string.Format("第{0:d4}段, {1} VIP= {2:X8} , HANDLE= {3:X8} , {4}", i, firstFound, cip, handle_addr, h_type.ToString());
            Console.WriteLine(text);


            if (!removeMore)
                removeMore = con.Codes.Any(o => o.EndPrefix);



            CollectHandler(cip, con, handler, _read_next, removeMore ? remove2 : remove1);

            con.Dump();

            ;

        }
    }

}
